package com.xh.hongaicodegenerator.langgraph4j.node;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.xh.hongaicodegenerator.langgraph4j.ai.ImageCollectionPlanService;
import com.xh.hongaicodegenerator.langgraph4j.ai.ImageCollectionPlanServiceFactory;
import com.xh.hongaicodegenerator.langgraph4j.model.ImageCollectionPlan;
import com.xh.hongaicodegenerator.langgraph4j.model.ImageResource;
import com.xh.hongaicodegenerator.langgraph4j.model.enums.ImageCategoryEnum;
import com.xh.hongaicodegenerator.langgraph4j.state.WorkflowContext;
import com.xh.hongaicodegenerator.langgraph4j.tools.ImageSearchTool;
import com.xh.hongaicodegenerator.langgraph4j.tools.LogoGeneratorTool;
import com.xh.hongaicodegenerator.langgraph4j.tools.MermaidDiagramTool;
import com.xh.hongaicodegenerator.langgraph4j.tools.UndrawIllustrationTool;
import com.xh.hongaicodegenerator.model.entity.User;
import com.xh.hongaicodegenerator.monitor.FluxContext;
import com.xh.hongaicodegenerator.monitor.MonitorContext;
import com.xh.hongaicodegenerator.monitor.MonitorContextHolder;
import com.xh.hongaicodegenerator.utils.SpringContextUtil;
import lombok.extern.slf4j.Slf4j;
import org.bsc.langgraph4j.action.AsyncNodeAction;
import org.bsc.langgraph4j.prebuilt.MessagesState;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;

/**
 * 图片收集节点（并发）
 */
@Slf4j
public class ImageCollectorNode {
    // 注入自定义线程池
    private static ExecutorService imageCollectionTaskExecutor;

    // 在适当位置初始化
    static {
        imageCollectionTaskExecutor = SpringContextUtil.getBean("imageCollectionTaskExecutor", ExecutorService.class);
    }

    public static AsyncNodeAction<MessagesState<String>> create() {
        return node_async(state -> {
            WorkflowContext context = WorkflowContext.getContext(state);
            String originalPrompt = context.getOriginalPrompt();
            List<ImageResource> collectedImages = new ArrayList<>();
            try {
                // 第一步：获取图片收集计划
                ImageCollectionPlanServiceFactory imageCollectionServiceFactory = SpringContextUtil.getBean(ImageCollectionPlanServiceFactory.class);
                ImageCollectionPlanService planService = imageCollectionServiceFactory.createImageCollectionPlanService();
                ImageCollectionPlan plan = planService.planImageCollection(originalPrompt);
                log.info("获取到图片收集计划，开始并发执行");

                // 第二步：并发执行各种图片收集任务
                List<CompletableFuture<List<ImageResource>>> futures = new ArrayList<>();

                // 并发执行内容图片搜索
                if (plan.getContentImageTasks() != null) {
                    ImageSearchTool imageSearchTool = SpringContextUtil.getBean(ImageSearchTool.class);
                    for (ImageCollectionPlan.ImageSearchTask task : plan.getContentImageTasks()) {
                        futures.add(CompletableFuture.supplyAsync(() -> {
                            try {
                                return imageSearchTool.searchContentImages(task.query());
                            } catch (Exception e) {
                                log.error("内容图片搜索任务失败, 查询词: {}", task.query(), e);
                                return new ArrayList<ImageResource>(); // 返回空列表而不是null
                            }
                        }, imageCollectionTaskExecutor));
                    }
                }

                // 并发执行插画图片搜索
                if (plan.getIllustrationTasks() != null) {
                    UndrawIllustrationTool illustrationTool = SpringContextUtil.getBean(UndrawIllustrationTool.class);
                    for (ImageCollectionPlan.IllustrationTask task : plan.getIllustrationTasks()) {
                        futures.add(CompletableFuture.supplyAsync(() -> {
                            try {
                                return illustrationTool.searchIllustrations(task.query());
                            } catch (Exception e) {
                                log.error("插画图片搜索任务失败, 查询词: {}", task.query(), e);
                                return new ArrayList<ImageResource>();
                            }
                        }, imageCollectionTaskExecutor));
                    }
                }

                // 并发执行架构图生成
                if (plan.getDiagramTasks() != null) {
                    MermaidDiagramTool diagramTool = SpringContextUtil.getBean(MermaidDiagramTool.class);
                    for (ImageCollectionPlan.DiagramTask task : plan.getDiagramTasks()) {
                        futures.add(CompletableFuture.supplyAsync(() -> {
                            try {
                                return diagramTool.generateMermaidDiagram(task.mermaidCode(), task.description());
                            } catch (Exception e) {
                                log.error("架构图生成任务失败, 描述: {}", task.description(), e);
                                return new ArrayList<ImageResource>();
                            }
                        }, imageCollectionTaskExecutor));
                    }
                }

                // 并发执行Logo生成
                if (plan.getLogoTasks() != null) {
                    LogoGeneratorTool logoTool = SpringContextUtil.getBean(LogoGeneratorTool.class);
                    for (ImageCollectionPlan.LogoTask task : plan.getLogoTasks()) {
                        futures.add(CompletableFuture.supplyAsync(() -> {
                            try {
                                return logoTool.generateLogos(task.description());
                            } catch (Exception e) {
                                log.error("Logo生成任务失败, 描述: {}", task.description(), e);
                                return new ArrayList<ImageResource>();
                            }
                        }, imageCollectionTaskExecutor));
                    }
                }

                // 等待所有任务完成（设置超时时间避免无限等待）
                try {
                    CompletableFuture<Void> allTasks = CompletableFuture.allOf(
                            futures.toArray(new CompletableFuture[0]));
                    // 设置超时时间，避免任务长时间挂起
                    allTasks.get(60, TimeUnit.SECONDS); // 30秒超时
                } catch (TimeoutException e) {
                    log.warn("图片收集任务超时，继续处理已完成的任务");
                } catch (Exception e) {
                    log.error("等待任务完成时发生错误: {}", e.getMessage(), e);
                }

                for (CompletableFuture<List<ImageResource>> future : futures) {
                    List<ImageResource> images = future.get();
                    if (images != null) {
                        collectedImages.addAll(images);
                    }
                }

            } catch (Exception e) {
                log.error("图片收集失败: {}", e.getMessage(), e);
            }
            // 更新状态
            context.setCurrentStep("图片收集");

            // 对收集到的图片进行数量限制：每种素材最多保留 10 个
            Map<ImageCategoryEnum, List<ImageResource>> grouped = collectedImages.stream()
                    .collect(Collectors.groupingBy(ImageResource::getCategory));

            List<ImageResource> limitedImages = new ArrayList<>();
            for (Map.Entry<ImageCategoryEnum, List<ImageResource>> entry : grouped.entrySet()) {
                List<ImageResource> images = entry.getValue();
                if (images.size() > 10) {
                    limitedImages.addAll(images.subList(0, 10)); // 只取前10个
                } else {
                    limitedImages.addAll(images);
                }
            }

            // 更新状态
            context.setImageList(limitedImages);
            Integer stepCount = context.getStepCount()+1;
            context.setStepCount(stepCount);
            String imageListStr = context.getImageListStr();
            // 构建增强后的提示词
            // 如果有图片资源，则添加图片信息
            StringBuilder imageListStrBuilder = new StringBuilder();
            if (CollUtil.isNotEmpty(limitedImages) || StrUtil.isNotBlank(imageListStr)) {
                imageListStrBuilder.append("\n\n利用工具搜集到的可用素材资源：\n\n");
                if (CollUtil.isNotEmpty(limitedImages)) {
                    for (ImageResource image : limitedImages) {
                        imageListStrBuilder.append("- ")
                                .append(image.getCategory().getText())
                                .append("：")
                                .append(image.getDescription())
                                .append("（")
                                .append(image.getUrl())
                                .append("）\n");
                    }
                } else {
                    imageListStrBuilder.append(imageListStr);
                }
            }
            context.setImageListStr(imageListStrBuilder.toString());
            log.info("并发图片收集完成，共收集到 {} 张图片", limitedImages.size());
            FluxContext flux= MonitorContextHolder.getFlux();
            flux.getSink().next(String.format("\n\n第 %s 步 %s 执行完成，共搜集到 %s 张图片\n\n",stepCount,"图片收集",limitedImages.size()));
            return WorkflowContext.saveContext(context);
        });
    }
}
